#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基础示例：DUCA算法演示
===================

本示例演示如何使用DUCA算法解决一个简单的分布式二次规划问题。

问题描述：
- 5个节点的网络
- 每个节点有一个二次目标函数
- 全局约束：所有节点决策变量之和小于等于某个值

作者：AI Assistant
日期：2024
"""

import numpy as np
import matplotlib.pyplot as plt
import sys
import os

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def main():
    """
    主函数：运行基础DUCA示例
    """
    print("=" * 50)
    print("DUCA算法基础示例")
    print("=" * 50)
    
    try:
        # 导入必要的模块
        from src.network import Network
        from src.problem import DistributedProblem
        from algorithms.duca import DUCA
        
        print("✓ 模块导入成功")
        
        # 1. 创建网络拓扑
        print("\n1. 创建网络拓扑...")
        num_nodes = 5
        network = Network(num_nodes=num_nodes, topology='complete')
        stats = network.get_network_statistics()
        print(f"   - 节点数量: {num_nodes}")
        print(f"   - 网络类型: 完全连接图")
        print(f"   - 连接边数: {stats['num_edges']}")
        print(f"   - 网络连通性: {'是' if stats['is_connected'] else '否'}")
        
        # 2. 定义分布式优化问题
        print("\n2. 定义优化问题...")
        problem = DistributedProblem(network)
        
        # 设置简单的二次目标函数
        def create_local_objective(node_id):
            """为每个节点创建局部目标函数"""
            def f_i(x):
                # 简单的二次函数: (x - node_id)^2
                return np.sum((x - node_id) ** 2)
            return f_i
        
        # 为每个节点设置目标函数
        for i in range(num_nodes):
            problem.set_local_objective(i, create_local_objective(i))
        
        # 设置全局约束：所有变量之和 <= 10
        def global_constraint(x_all):
            return np.sum(x_all) - 10.0
        
        problem.set_global_constraint(global_constraint)
        print("   - 局部目标函数: 二次函数")
        print("   - 全局约束: 所有变量之和 ≤ 10")
        
        # 3. 创建DUCA算法实例
        print("\n3. 初始化DUCA算法...")
        algorithm = DUCA(
            problem=problem, 
            network=network,
            rho=1.0,  # 惩罚参数
            max_iterations=100,
            tolerance=1e-6
        )
        print("   - 算法类型: DUCA")
        print("   - 惩罚参数 ρ: 1.0")
        print("   - 最大迭代次数: 100")
        print("   - 收敛容差: 1e-6")
        
        # 4. 运行算法
        print("\n4. 运行算法...")
        print("   开始优化...")
        
        result = algorithm.solve(verbose=True)
        
        # 5. 显示结果
        print("\n5. 优化结果:")
        print("=" * 30)
        print(f"收敛状态: {'成功' if result['converged'] else '未收敛'}")
        print(f"迭代次数: {result['iterations']}")
        print(f"最优目标值: {result['objective_value']:.6f}")
        print(f"约束违反度: {result['constraint_violation']:.6f}")
        
        print("\n各节点最优解:")
        for i, x_opt in enumerate(result['solution']):
            print(f"   节点 {i+1}: x = {x_opt:.4f}")
        
        print(f"\n全局约束检查: Σx = {np.sum(result['solution']):.4f} ≤ 10")
        
        # 6. 绘制收敛曲线
        if 'history' in result:
            print("\n6. 绘制收敛曲线...")
            plot_convergence(result['history'])
            
        print("\n" + "=" * 50)
        print("示例运行完成！")
        print("=" * 50)
        
    except ImportError as e:
        print(f"❌ 模块导入失败: {e}")
        print("请确保已安装所有依赖包并激活虚拟环境")
        return False
    except Exception as e:
        print(f"❌ 运行出错: {e}")
        return False
    
    return True

def plot_convergence(history):
    """
    绘制算法收敛曲线
    
    参数:
        history: dict, 包含算法运行历史数据
    """
    try:
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
        
        # 目标函数值收敛曲线
        if 'objective_values' in history:
            ax1.plot(history['objective_values'], 'b-', linewidth=2)
            ax1.set_xlabel('迭代次数')
            ax1.set_ylabel('目标函数值')
            ax1.set_title('目标函数收敛曲线')
            ax1.grid(True, alpha=0.3)
        
        # 约束违反度收敛曲线
        if 'constraint_violations' in history:
            ax2.semilogy(history['constraint_violations'], 'r-', linewidth=2)
            ax2.set_xlabel('迭代次数')
            ax2.set_ylabel('约束违反度 (对数尺度)')
            ax2.set_title('约束违反度收敛曲线')
            ax2.grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig('duca_convergence.png', dpi=300, bbox_inches='tight')
        plt.show()
        print("   - 收敛曲线已保存为 'duca_convergence.png'")
        
    except Exception as e:
        print(f"   - 绘图失败: {e}")

if __name__ == "__main__":
    success = main()
    if not success:
        sys.exit(1) 